#include <iostream>
#include <stdio.h>
// #include <gtest/gtest.h>
#include <hip/amd_detail/hip_prof_str.h>
#include <hip/hip_runtime.h>
#include <roctracer.h>
#include <roctracer_ext.h>
#include <unistd.h>
// #include <hip_ostream_ops.h>

#define NSTREAM 2
#define BDIM 512

static int api_callback_conter = 0;
static int pool_active_cb_counter = 0;

// Macro to check ROC-tracer calls status
#define ROCTRACER_CHECK(error)                                                 \
  {                                                                            \
    roctracer_status_t localError = error;                                     \
    if (localError != ROCTRACER_STATUS_SUCCESS) {                              \
      std::cerr << "Error: " << roctracer_error_string()                       \
                << "\n    Code: " << localError << "\n    Str: " << #error     \
                << "\n    In File: " << __FILE__                               \
                << "\n    At line: " << __LINE__ << std::endl;                 \
      exit(-1);                                                                \
    }                                                                          \
  }

#define HIP_CHECK(error)                                                       \
  {                                                                            \
    hipError_t localError = error;                                             \
    if (localError != hipSuccess) {                                            \
      std::cerr << "Error: " << hipGetErrorString(localError)                  \
                << "\n    Code: " << localError << "\n    Str: " << #error     \
                << "\n    In File: " << __FILE__                               \
                << "\n    At line: " << __LINE__ << std::endl;                 \
      exit(-1);                                                                \
    }                                                                          \
  }

static void initialData(float *ip, int size) {
  int i;

  for (i = 0; i < size; i++) {
    ip[i] = (float)(rand() & 0xFF) / 10.0f;
    // printf("%f\n", ip[i]);
  }
}

__launch_bounds__(1024)
__global__ void sumArrays(float *A, float *B, float *C, const int N) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;

  if (idx < N) {
    for (int j = 0; j < 60; j++) {
      C[idx] = A[idx] + B[idx];
    }
  }
}

static bool checkResult(float *hostRef, float *gpuRef, const int N) {
  double epsilon = 1.0E-8;
  bool match = 1;

  for (int i = 0; i < N; i++) {
    if (abs(hostRef[i] - gpuRef[i]) > epsilon) {
      match = 0;
      printf("Arrays do not match!\n");
      printf("host %5.2f gpu %5.2f at %d\n", hostRef[i], gpuRef[i], i);
      return false;
    }
  }

  if (match)
    printf("Arrays match.\n\n");
  return true;
}

void activity_rtapi_callback(uint32_t domain, uint32_t cid,
                             const void *callback_data, void *arg) {
  const hip_api_data_t *data =
      reinterpret_cast<const hip_api_data_t *>(callback_data);
  uint64_t timestamp;
  roctracer_get_timestamp(&timestamp);
  const char *op_name = hip_api_name(cid);
  if (data->phase == ACTIVITY_API_PHASE_ENTER) {
    std::cout << "test roctracer >> enter api " << op_name << " time is "
              << timestamp << std::endl;
  } else {
    std::cout << "test roctracer >> return api " << op_name << " time is "
              << timestamp << std::endl;
  }
  api_callback_conter++;
}

#define MEMORY_POOL_BUFFER_SIZE 0x80000

// Activity tracing callback
void mempool_callback(const char *begin, const char *end, void *arg) {
  const roctracer_record_t *record =
      reinterpret_cast<const roctracer_record_t *>(begin);
  const roctracer_record_t *end_record =
      reinterpret_cast<const roctracer_record_t *>(end);
  printf("###mempool_callback\n");
  while (record < end_record) {
    const char *name =
        roctracer_op_string(record->domain, record->op, record->kind);
    printf(
        "pool_activity_callback(\"%s\"): domain(%u) op(%u) kind(%u) record(%p) "
        "correlation_id(%lu)\n",
        name, record->domain, record->op, record->kind, record,
        record->correlation_id);
    roctracer_next_record(record, &record);
    pool_active_cb_counter++;
  }
  printf("###mempool_callback finish\n");
}

int main() {
  api_callback_conter = 0;
  pool_active_cb_counter = 0;
  hipError_t hip_ret = hipSuccess;
  roctracer_status_t trac_ret = ROCTRACER_STATUS_SUCCESS;
  int dev = 0;
  hip_ret = hipSetDevice(dev);
  HIP_CHECK(hip_ret);
  uint64_t starttime;
  trac_ret = roctracer_get_timestamp(&starttime);
  ROCTRACER_CHECK(trac_ret);
  std::cout << "test roctracer >> start time is " << starttime << std::endl;
  trac_ret = roctracer_enable_op_callback(ACTIVITY_DOMAIN_HIP_API, 1,
                                          activity_rtapi_callback, nullptr);
  ROCTRACER_CHECK(trac_ret);
  trac_ret = roctracer_enable_domain_callback(ACTIVITY_DOMAIN_HIP_API,
                                              activity_rtapi_callback, nullptr);
  ROCTRACER_CHECK(trac_ret);
  trac_ret = roctracer_disable_op_callback(ACTIVITY_DOMAIN_HIP_API, 1);
  ROCTRACER_CHECK(trac_ret);
  trac_ret = roctracer_disable_domain_callback(ACTIVITY_DOMAIN_HIP_API);
  ROCTRACER_CHECK(trac_ret);
  trac_ret = roctracer_disable_callback();
  ROCTRACER_CHECK(trac_ret);
  trac_ret = roctracer_enable_callback(activity_rtapi_callback, nullptr);
  ROCTRACER_CHECK(trac_ret);
  roctracer_properties_t properties{};
  properties.buffer_size = MEMORY_POOL_BUFFER_SIZE;
  properties.buffer_callback_fun = mempool_callback;
  trac_ret = roctracer_open_pool(&properties);
  ROCTRACER_CHECK(trac_ret);
  trac_ret = roctracer_disable_activity();
  ROCTRACER_CHECK(trac_ret);
  trac_ret = roctracer_enable_activity();
  ROCTRACER_CHECK(trac_ret);
  hipDeviceProp_t deviceProp;
  hip_ret = hipGetDeviceProperties(&deviceProp, dev);
  HIP_CHECK(hip_ret);
  printf("> Using Device %d: %s\n", dev, deviceProp.name);

  // set up data size of vectors
  // int nElem = 1 << 2;
  int nElem = 1 << 8;
  printf("> vector size = %d\n", nElem);
  size_t nBytes = nElem * sizeof(float);
  printf("> size nBytes = %ld\n", nBytes);

  // malloc pinned host memory for async memcpy
  float *h_A, *h_B, *h_C;
  hip_ret = hipHostMalloc((void **)&h_A, nBytes, hipHostMallocDefault);
  HIP_CHECK(hip_ret);
  hip_ret = hipHostMalloc((void **)&h_B, nBytes, hipHostMallocDefault);
  HIP_CHECK(hip_ret);
  hip_ret = hipHostMalloc((void **)&h_C, nBytes, hipHostMallocDefault);
  HIP_CHECK(hip_ret);

  // initialize data at host side
  initialData(h_A, nElem);
  initialData(h_B, nElem);
  memset(h_C, 0, nBytes);

  // add vector at host side for result checks
  // sumArraysOnHost(h_A, h_B, hostRef, nElem);

  // malloc device global memory
  float *d_A, *d_B, *d_C;
  hip_ret = hipMalloc((float **)&d_A, nBytes);
  HIP_CHECK(hip_ret);
  hip_ret = hipMalloc((float **)&d_B, nBytes);
  HIP_CHECK(hip_ret);
  hip_ret = hipMalloc((float **)&d_C, nBytes);
  HIP_CHECK(hip_ret);

  hipEvent_t start, stop;
  hip_ret = hipEventCreate(&start);
  HIP_CHECK(hip_ret);
  hip_ret = hipEventCreate(&stop);
  HIP_CHECK(hip_ret);

  // invoke kernel at host side
  dim3 block(BDIM);
  dim3 grid((nElem + block.x - 1) / block.x);
  printf("> grid (%d,%d) block (%d,%d)\n", grid.x, grid.y, block.x, block.y);

  // warmup
  roctracer_stop();
  hip_ret = hipMemcpy(d_A, h_A, nBytes, hipMemcpyHostToDevice);
  HIP_CHECK(hip_ret);
  hip_ret = hipMemcpy(d_B, h_B, nBytes, hipMemcpyHostToDevice);
  HIP_CHECK(hip_ret);
  hipLaunchKernelGGL(sumArrays, dim3(grid), dim3(block), 0, 0, d_A, d_B, d_C,
                     nElem);
  hip_ret = hipMemcpy(h_C, d_C, nBytes, hipMemcpyDeviceToHost);
  HIP_CHECK(hip_ret);
  roctracer_start();
  // sequential operation
  hip_ret = hipEventRecord(start, 0);
  HIP_CHECK(hip_ret);
  hip_ret = hipMemcpy(d_A, h_A, nBytes, hipMemcpyHostToDevice);
  HIP_CHECK(hip_ret);
  hip_ret = hipMemcpy(d_B, h_B, nBytes, hipMemcpyHostToDevice);
  HIP_CHECK(hip_ret);
  hip_ret = hipEventRecord(stop, 0);
  HIP_CHECK(hip_ret);
  hip_ret = hipEventSynchronize(stop);
  HIP_CHECK(hip_ret);
  float memcpy_h2d_time;
  hip_ret = hipEventElapsedTime(&memcpy_h2d_time, start, stop);
  HIP_CHECK(hip_ret);

  hip_ret = hipEventRecord(start, 0);
  HIP_CHECK(hip_ret);
  hipLaunchKernelGGL(sumArrays, dim3(grid), dim3(block), 0, 0, d_A, d_B, d_C,
                     nElem);
  hip_ret = hipEventRecord(stop, 0);
  HIP_CHECK(hip_ret);
  hip_ret = hipEventSynchronize(stop);
  HIP_CHECK(hip_ret);
  float kernel_time;
  hip_ret = hipEventElapsedTime(&kernel_time, start, stop);
  HIP_CHECK(hip_ret);

  hip_ret = hipEventRecord(start, 0);
  HIP_CHECK(hip_ret);
  hip_ret = hipMemcpy(h_C, d_C, nBytes, hipMemcpyDeviceToHost);
  HIP_CHECK(hip_ret);
  hip_ret = hipEventRecord(stop, 0);
  HIP_CHECK(hip_ret);
  hip_ret = hipEventSynchronize(stop);
  HIP_CHECK(hip_ret);
  float memcpy_d2h_time;
  hip_ret = hipEventElapsedTime(&memcpy_d2h_time, start, stop);
  HIP_CHECK(hip_ret);

  printf("Measured timings (throughput):\n");
  printf(" Memcpy host to device\t: %f ms (%f GB/s)\n", memcpy_h2d_time,
         (2 * nBytes * 1e-6) / memcpy_h2d_time);
  printf(" Memcpy device to host\t: %f ms (%f GB/s)\n", memcpy_d2h_time,
         (nBytes * 1e-6) / memcpy_d2h_time);
  printf(" Kernel time: %f ms\n", kernel_time);
  float total_time = memcpy_h2d_time + memcpy_d2h_time + kernel_time;
  printf(" Total time: %f ms\n", total_time);

  // grid parallel operation
  int iElem = nElem / NSTREAM;
  size_t iBytes = iElem * sizeof(float);
  printf("> iElem = %d\n", iElem);
  printf("> iBytes = %ld\n", iBytes);
  grid.x = (iElem + block.x - 1) / block.x;
  printf("> grid (%d,%d) block (%d,%d)\n", grid.x, grid.y, block.x, block.y);

  hipStream_t stream[NSTREAM];
  for (int i = 0; i < NSTREAM; ++i) {
    hip_ret = hipStreamCreate(&stream[i]);
    HIP_CHECK(hip_ret);
  }

  hip_ret = hipEventRecord(start, 0);
  HIP_CHECK(hip_ret);
  // initiate all asynchronous transfers to the device
  for (int i = 0; i < NSTREAM; ++i) {
    int ioffset = i * iElem;
    hip_ret = hipMemcpyAsync(&d_A[ioffset], &h_A[ioffset], iBytes,
                             hipMemcpyHostToDevice, stream[i]);
    HIP_CHECK(hip_ret);
    hip_ret = hipMemcpyAsync(&d_B[ioffset], &h_B[ioffset], iBytes,
                             hipMemcpyHostToDevice, stream[i]);
    HIP_CHECK(hip_ret);
  }

  // launch a kernel in each stream
  for (int i = 0; i < NSTREAM; ++i) {
    int ioffset = i * iElem;
    hipLaunchKernelGGL(sumArrays, dim3(grid), dim3(block), 0, stream[i],
                       &d_A[ioffset], &d_B[ioffset], &d_C[ioffset], iElem);
  }

  // enqueue asynchronous transfers from the device
  for (int i = 0; i < NSTREAM; ++i) {
    int ioffset = i * iElem;
    hip_ret = hipMemcpyAsync(&h_C[ioffset], &d_C[ioffset], iBytes,
                             hipMemcpyDeviceToHost, stream[i]);
    HIP_CHECK(hip_ret);
  }

  hip_ret = hipEventRecord(stop, 0);
  HIP_CHECK(hip_ret);
  hip_ret = hipEventSynchronize(stop);
  HIP_CHECK(hip_ret);
  float execution_time;
  hip_ret = hipEventElapsedTime(&execution_time, start, stop);
  HIP_CHECK(hip_ret);

  printf("Actual results from overlapped data transfers:\n");
  printf("overlap with %d streams : %f ms\n", NSTREAM, execution_time);

  // check device results
  // EXPECT_TRUE(checkResult(hostRef, gpuRef, nElem));

  // free device global memory
  hip_ret = hipFree(d_A);
  HIP_CHECK(hip_ret);
  hip_ret = hipFree(d_B);
  HIP_CHECK(hip_ret);
  hip_ret = hipFree(d_C);
  HIP_CHECK(hip_ret);

  // free host memory
  hip_ret = hipHostFree(h_A);
  HIP_CHECK(hip_ret);
  hip_ret = hipHostFree(h_B);
  HIP_CHECK(hip_ret);
  hip_ret = hipHostFree(h_C);
  HIP_CHECK(hip_ret);

  // destroy events
  hip_ret = hipEventDestroy(start);
  HIP_CHECK(hip_ret);
  hip_ret = hipEventDestroy(stop);
  HIP_CHECK(hip_ret);

  // destroy streams
  for (int i = 0; i < NSTREAM; ++i) {
    hip_ret = hipStreamDestroy(stream[i]);
    HIP_CHECK(hip_ret);
  }

  hip_ret = hipDeviceReset();
  HIP_CHECK(hip_ret);
  trac_ret = roctracer_flush_activity();
  ROCTRACER_CHECK(trac_ret);
  roctracer_stop();
  printf("test roctracer api callback counter is %d check line is 110\n",
         api_callback_conter);
  // EXPECT_TRUE(api_callback_conter > 110); //api callback
  printf("test roctracer api active counter is %d check line is 60\n",
         pool_active_cb_counter);
  // EXPECT_TRUE(pool_active_cb_counter > 60); //active callback
  return 0;
}
